#include "attack.h"
#include <gmp.h>
#include <math.h>
#include <openssl/sha.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

/* The lHash for the OAEP-decode */
Byte lHash[] = { 0xDA , 0x39 , 0xA3 , 0xEE , \
                 0x5E , 0x6B , 0x4B , 0x0D , \
                 0x32 , 0x55 , 0xBF , 0xEF , \
                 0x95 , 0x60 , 0x18 , 0x90 , \
                 0xAF , 0xD8 , 0x07 , 0x09 , \
               } ;

int main(int argc, char *argv[])
{
  Options opt = parse_args(argc, argv);
  Challenge challenge = parse_challenge(opt.challenge_path);


  mpz_t f1, f2;
  mpz_t b;
  mpz_t m;

  size_t k;
  char *oaep_message;
  char *message;

  //char *message;
  mpz_init(f1);
  mpz_init(f2);
  mpz_init(b);
  mpz_init(m);

  // Calculate the value of B
  calculate_b(b, challenge);

  
  // THE ATTACK!
  step_1(f1,           opt, challenge);
  step_2(f1, f2, b,    opt, challenge);
  step_3(    f2, b, m, opt, challenge);

  // THE DECRYPTION!
  // mpz_sizeinbase only supports bases between 2 and 62 (but it works anyway)
  // so to find the size-in-base 256 we take halve the size-in-base 16 and take
  // the ceiling.  ...anything to please Valgrind eh?
  //k = mpz_sizeinbase(challenge.n, 256);
  k = mpz_sizeinbase(challenge.n, 16);
  if (2*(k/2) < k) k = k/2 +1;
  else             k = k/2;

  oaep_message = i2osp(m, k);
  message = eme_oaep_decode(oaep_message, k);

  printf("%s\n", message);

  mpz_clear(challenge.n);
  mpz_clear(challenge.e);
  mpz_clear(challenge.c);
  free(opt.program_name);
  free(opt.oracle_path);
  free(opt.challenge_path);
  free(message);
  free(oaep_message);

  mpz_clear(f1);
  mpz_clear(f2);
  mpz_clear(b);
  mpz_clear(m);


  return EXIT_SUCCESS;
}

Options parse_args(int c, char *v[])
{
  Options opt;
  if (c != 3)
  {
    usage(v[0]);
    exit(EXIT_FAILURE);
  }

  opt.program_name   = strdup(v[0]);
  opt.oracle_path    = strdup(v[1]);
  opt.challenge_path = strdup(v[2]);

  return opt;
}

void usage(char *name)
{
  fprintf(stderr, "USAGE: %s oracle challenge\n", name);
}

Challenge parse_challenge(char *path)
{
  Challenge in;
  FILE *fp; 

  if ((fp = fopen(path, "r")) == NULL)
  {
    perror("challenge did not exist or wasn't readable");
    exit(EXIT_FAILURE);
  }

  mpz_init(in.n);
  mpz_init(in.e);
  mpz_init(in.c);

  mpz_inp_str(in.n, fp, 16);
  mpz_inp_str(in.e, fp, 16);
  mpz_inp_str(in.c, fp, 16);
  
  fclose(fp);

  gmp_fprintf(stderr, "Read in challenge.\n"
                      "  n: %ZX\n"
                      "  e: %ZX\n"
                      "  c: %ZX\n\n",
                      in.n, in.e, in.c);

  return(in);
}

int send_to_oracle(char *oracle, Challenge challenge)
{
  char *file = "jh6970.temp";
  char *null = "/dev/null";

  FILE *np;
  FILE *fp;
  char command[strlen(oracle)+2+strlen(file)+1];
  int exit_code;

  fp = fopen(file, "w");
  np = fopen(null, "w");

  // Make sure we print c as an octet string
  // WARNING: Hardcoded c as 128B
  gmp_fprintf(fp, "%ZX\n%ZX\n%0256ZX\n", 
      challenge.n, challenge.e, challenge.c);

  /*
   *printf("send_to_oracleING:\n");
   *gmp_printf("%ZX\n%ZX\n%0256ZX\n", 
   *    challenge.n, challenge.e, challenge.c);
   */

  fclose(fp);
  fclose(np);

  sprintf(command, "./%s <%s", oracle, file);
  //printf("%s\n", command);
  exit_code = system(command);
  exit_code = exit_code>>8; 
  
  switch(exit_code)
  {
    // Expected error codes
    case 0: case 1: case 2:
      return exit_code;

    case 4:
      gmp_fprintf(stderr, 
          "Error #4: RSADP failed because the operand was out of range.\n"
          "i.e. the ciphertext is not in [0..N-1).\n"
          "  n: 0x%ZX\n"
          "  e: 0x%ZX\n"
          "  c: 0x%0256ZX\n",
          challenge.n, challenge.e, challenge.c
        );
      exit(EXIT_FAILURE);

    case 6:
      gmp_fprintf(stderr, 
          "Error #6: RSAES-OAEP-DECRYPT failed because of a length check.\n"
          "i.e. ciphertext doesn't match the length of N.\n"
          "  n: 0x%0ZX\n"
          "  e: 0x%0ZX\n"
          "  c: 0x%0256ZX\n",
          challenge.n, challenge.e, challenge.c
        );
      exit(EXIT_FAILURE);

    case 7:
      gmp_fprintf(stderr, 
          "Error #7: RSAES-DECRYPT failed because of a length check.\n"
          "i.e. ciphertext doesn't match the length of the hash function output.\n"
          "  n: 0x%0ZX\n"
          "  e: 0x%0ZX\n"
          "  c: 0x%0256ZX\n",
          challenge.n, challenge.e, challenge.c
        );
      exit(EXIT_FAILURE);

    default:
      gmp_fprintf(stderr,
          "Error #%d: Unknown error\n"
          "  n: 0x%0ZX\n"
          "  e: 0x%0ZX\n"
          "  c: 0x%0256ZX\n",
          exit_code,
          challenge.n, challenge.e, challenge.c
        );
      exit(EXIT_FAILURE);
  }
}

void calculate_b(mpz_t b, Challenge challenge)
{
  mpz_t temp;
  unsigned long int k = mpz_sizeinbase(challenge.n, 256);

  mpz_set_ui(b, 2);
  mpz_powm_ui(b, b, 8*(k-1), challenge.n);

  // The attack assumes that 2*b < n: but is it?
  mpz_init(temp);
  mpz_mul_ui(temp, b, 2UL);
  if(mpz_cmp(temp, challenge.n) >= 0)
  {
    gmp_fprintf(stderr, "Error: 2*B is greater than or equal to N.\n"
                        "See section 3.2 of Manger's paper.\n"
                        "  k: %d\n"
                        "  b: 0x%ZX\n"
                        "  n: 0x%ZX\n",
                        k, b, challenge.n);

    exit(EXIT_FAILURE);
  }

  gmp_fprintf(stderr, "Finished calclulating B:\n"
                      "  b: %ZX\n"
                      "  k: %d\n\n",
                      b,k);
  mpz_clear(temp); 
}


void step_1(mpz_t f1, Options opt, Challenge challenge)
{
  Challenge test;
  int exit_code;
  int count = 0;

  mpz_init(test.n);
  mpz_init(test.e);
  mpz_init(test.c);
  
  mpz_set(test.n, challenge.n);
  mpz_set(test.e, challenge.e);

  gmp_fprintf(stderr, "Starting Step #1\n");

  // Step 1.1 
  mpz_set_ui(f1, 2UL);
  do 
  {
    count++;
    // Step 1.2
    mpz_powm(test.c, f1, challenge.e, challenge.n);
    mpz_mul(test.c, test.c, challenge.c);
    mpz_mod(test.c, test.c, challenge.n);

    exit_code = send_to_oracle(opt.oracle_path, test);

    // Step 1.3a
    if (exit_code != 1)
      mpz_mul_ui(f1, f1, 2UL);
  }
  while (exit_code != 1);
  // Step 1.3b

  mpz_clear(test.n);
  mpz_clear(test.e); 
  mpz_clear(test.c); 

  gmp_fprintf(stderr, "Finished Step #1 after %d iterations.\n"
                      "  f1: %ZX\n\n",
                      count, f1);
}

/* Note:
 *   The attack states that as f2 increases at iterations through step
 *   2.3a the lower bound on f2.m increases eventuall exceeding n when
 *   f2 = ceil(2n/b). Branch 2.3b must occur at or before this multiple.
 *
 *   With the values of N, and B we have in the challenge there is an
 *   upper bound of 97 iterations.
 */
void step_2(mpz_t f1, mpz_t f2, mpz_t b, Options opt, Challenge challenge)
{
  Challenge test;
  int exit_code;
  int count = 0;
  mpz_t f1d2;
  mpz_t limit;

  mpz_init(test.n);
  mpz_init(test.e);
  mpz_init(test.c);
  mpz_init(f1d2);
  mpz_init(limit);

  mpz_set(test.n, challenge.n);
  mpz_set(test.e, challenge.e);
  mpz_fdiv_q_ui(f1d2, f1, 2UL);

  // Calculate the upper bound on the number of iterations
  mpz_set(limit, test.n);
  mpz_mul_ui(limit, limit, 2UL);
  mpz_cdiv_q(limit, limit, b);
  gmp_fprintf(stderr, "Starting Step #2.\n" 
                      "  This should take no more than %Zd iterations\n\n",
                      limit);

  // Step 2.1:
  mpz_add(f2, challenge.n, b);
  mpz_fdiv_q(f2, f2, b);
  mpz_mul(f2, f2, f1d2);

  do 
  {
    count++;
    if (mpz_cmp_ui(limit, count) < 0)
    {
      gmp_fprintf(stderr, "Error: Step #2 is taking too long to complete\n"
                          "  limit: %Zd\n"
                          "  count: %d\n",
                          limit, count);
    }

    // Step 2.2:
    mpz_powm(test.c, f2, test.e, test.n);
    mpz_mul(test.c, test.c, challenge.c);
    mpz_mod(test.c, test.c, test.n);
    //mpz_mod(test.c, f2, test.n);
    exit_code = send_to_oracle(opt.oracle_path, test);

    // Step 2.3a:
    if (exit_code == 1)
      mpz_add(f2, f2, f1d2); 
    
  } 
  while(exit_code == 1);
  // Step 2.3b:

  gmp_fprintf(stderr, "Finished Step #2 after %d iterations.\n"
                      "  f2: %ZX\n\n",
                      count, f2);


  mpz_clear(f1d2);
  mpz_clear(limit);
  mpz_clear(test.n);
  mpz_clear(test.e);
  mpz_clear(test.c);
}

void step_3(mpz_t f2, mpz_t b, mpz_t m, Options opt, Challenge challenge)
{
  int exit_code;
  int count = 0;
  //int x, y;
  mpz_t ftmp;
  mpz_t f3;
  mpz_t mmin;
  mpz_t mmax;
  mpz_t i;
  mpz_t in;
  mpz_t tmp;
  Challenge test;

  mpz_init(ftmp);
  mpz_init(f3);
  mpz_init(mmin);
  mpz_init(mmax);
  mpz_init(i);
  mpz_init(in);
  mpz_init(tmp);
  mpz_init(test.n);
  mpz_init(test.e);
  mpz_init(test.c);
  
  mpz_set(test.n, challenge.n);
  mpz_set(test.e, challenge.e);

  // Step 3.1:
  mpz_cdiv_q(mmin, challenge.n, f2);
  mpz_add(mmax, challenge.n, b);
  mpz_fdiv_q(mmax, mmax, f2);   

  gmp_fprintf(stderr, "Starting Step #3\n");
                   
  /*
  // Sanity check: f2(mmax-mmin) ~ B                                  
  mpz_sub(tmp, mmax, mmin);                                           
  mpz_mul(tmp, f2, tmp);                                              
  mpz_sub(tmp, b, tmp);                                               
  gmp_fprintf(stderr, "In Step #3 f2*(mmax-mmin) differs from b by:\n"
                      "  (base 10) %Zd\n\n",                          
                      tmp);                                           
  */

  do
  {
    count++;

    // Step 3.2:
    mpz_mul_ui(ftmp, b, 2UL);
    mpz_sub(tmp, mmax, mmin);
    mpz_fdiv_q(ftmp, ftmp, tmp);

    /*
    // Sanity check: width(ftmp*m) ~ width(2B)
    // QUESTION: What does width mean?
    // GUESS: width in chars?
    mpz_mul(tmp, ftmp, mmin);
    x = mpz_sizeinbase(tmp, 256);
    mpz_mul_ui(tmp, b, 2UL);
    y = mpz_sizeinbase(tmp, 256);
    */
    /*
    gmp_fprintf(stderr, "On iteration %d of Step #3 (ftmp * mmin) differs from 2b in width by:\n"
                        "  width(ftmp*mmin): %d\n"
                        "  width(2b):        %d\n"
                        "  difference:       %d\n\n",
                        count, x, y, x-y);
    */

    // Step 3.3
    mpz_mul(i, ftmp, mmin);
    mpz_fdiv_q(i, i, test.n);
    mpz_mul(in, i, test.n);



    // Step 3.4
    mpz_cdiv_q(f3, in, mmin);
    mpz_powm(test.c, f3, test.e, test.n);
    mpz_mul(test.c, test.c, challenge.c);
    mpz_mod(test.c, test.c, test.n);

    gmp_fprintf(stderr, "On iteration %d of Step #3.\n"
                        "  mmin: %ZX\n"
                        "  mmax: %ZX\n"
                        "  f3:   %ZX\n"
                        "  ftmp: %ZX\n"
                        "  i:    %ZX\n"
                        "  in:   %ZX\n\n",
                        count, mmin, mmax, f3, ftmp, i, in); 

    exit_code = send_to_oracle(opt.oracle_path, test);
    // printf("I GOT EXIT CODE: %d\n", exit_code);


    // Step 3.5a
    if (exit_code == 1)
    {
      mpz_add(mmin, in, b);
      mpz_cdiv_q(mmin, mmin, f3);
    } // Step 3.5b
    else if (exit_code == 2)
    {
      mpz_add(mmax, in, b);
      mpz_fdiv_q(mmax, mmax, f3);
    }

  } while (mpz_cmp(mmin, mmax) != 0);

  mpz_set(m, mmin);

  gmp_fprintf(stderr, "Finished Step #3 after %d iterations.\n"
                      "  f3:   %ZX\n"
                      "  m:    %ZX\n"
                      "  in:   %ZX\n\n",
                      count, f3, m, in);

  mpz_clear(ftmp);
  mpz_clear(f3);
  mpz_clear(mmin);
  mpz_clear(mmax);
  mpz_clear(i);
  mpz_clear(in);
  mpz_clear(tmp);
  mpz_clear(test.n);
  mpz_clear(test.e);
  mpz_clear(test.c);
}


/* Convert a hex char to an int */
int h2int(char x)
{
  switch (x)
  {
    case '0':           return 0;
    case '1':           return 1;
    case '2':           return 2;
    case '3':           return 3;
    case '4':           return 4;
    case '5':           return 5;
    case '6':           return 6;
    case '7':           return 7;
    case '8':           return 8;
    case '9':           return 9;
    case 'A': case 'a': return 10;
    case 'B': case 'b': return 11;
    case 'C': case 'c': return 12;
    case 'D': case 'd': return 13;
    case 'E': case 'e': return 14;
    case 'F': case 'f': return 15;

    default: 
      fprintf(stderr, "Error: h2int: Invalid character '%c' in hex digit\n", x);
      exit(EXIT_FAILURE);
  }
}


/* Translate the message to an octet string */
char *i2osp(mpz_t m, int k)
{
  char hex[2*k];
  char *str = calloc(k, sizeof(*str));

  gmp_sprintf(hex, "%0256ZX", m);

  for (int i = 0; i < k; i++)
    str[i] = (h2int(hex[2*i]) << 4) + h2int(hex[2*i+1]);

  return str;
}


char *eme_oaep_decode(char *EM, size_t k)
{
  Byte Y;
  Byte maskedSeed[hLen];
  Byte maskedDB[k-hLen-1];

  Byte *seedMask;
  Byte seed[hLen];

  Byte *dbMask;
  Byte DB[k-hLen-1];

  Byte lHash_check[hLen];

  char *M;

  size_t n;

  fprintf(stderr, "Starting EME-OAEP expansion\n");


  // Error Check: k > 2*hLen+2
  if (k < 2*hLen + 2)
  {
    fprintf(stderr, "Error: k < 2*hLen + 2\n");
    exit(EXIT_FAILURE);
  }

  // 3a
  // Hard coded this. No need to calculate
  //SHA1(L, strlen(L), lHash);

  // 3b
  Y = EM[0];
  for (size_t i=0; i<hLen; i++)
   maskedSeed[i] = EM[1+i];

  for (size_t i=0; i<k-hLen-1; i++)
   maskedDB[i] = EM[1+hLen+i];

  // 3c
  seedMask = MGF1(maskedDB, k-hLen-1, hLen);

  // 3d
  for (size_t i = 0; i < hLen; i++)
    seed[i] = maskedSeed[i] ^ seedMask[i];


  // 3e
  dbMask = MGF1(seed, hLen, k-hLen-1);

  // 3f
  for (size_t i = 0; i < k-hLen-1; i++)
    DB[i] = maskedDB[i] ^ dbMask[i];

  // 3g
  if (Y != 0)
  {
    fprintf(stderr, "Error: Y =/= 0\n");
    exit(EXIT_FAILURE);
  }

  memcpy(lHash_check, DB, hLen);
  if(memcmp(lHash, lHash_check, hLen) != 0)
  {
    fprintf(stderr, "Error: hash mismatch\n");
    exit(EXIT_FAILURE);
  }

  for (n = 0; DB[hLen+n] == 0; n++);
  if (DB[hLen + n] != 0x01)
  {
    fprintf(stderr, "Error: missing padding seperator\n");
    fprintf(stderr, "  DB[hLen + %lu] = %02x\n", n, DB[hLen + n]); 
    exit(EXIT_FAILURE);
  }

  M = calloc(k-2*hLen-n-1, sizeof(*M));
  for (size_t i = 0; i < k-2*hLen-n-2; i++)
    M[i] = DB[hLen +n +1 +i];
  M[k-2*hLen-n-2] = '\0';

  free(seedMask);
  free(dbMask);

  fprintf(stderr, "Finished EME-OAEP expansion\n");
  return M;
}

void print_bytes(Byte *bytes, const size_t len)
{
  for (size_t i = 0; i < len; i++)
  {
    fprintf(stderr, "%02x", bytes[i]);
    if ((i+1) % 4 == 0) putchar(' ');
  }
  fprintf(stderr, "\n");
}


Byte *MGF1(Byte *seed, size_t seedLen, size_t maskLen) 
{
  size_t max; 
  size_t dataLen;
  Byte *data; 
  Byte *output;
  Byte *tmp;

  dataLen = seedLen+4;
  data = calloc(dataLen, sizeof(*data));

  if(data == NULL) 
    return NULL;
  else 
    memcpy(data, seed, seedLen);

  if (maskLen % hLen == 0)
    max = maskLen / hLen;
  else
    max = (maskLen - (maskLen % hLen)) / hLen + 1;

  output = calloc(max*hLen, sizeof(*output));
  
  tmp = output;

  /* WARNING THIS CODE DEPENDS ON RUNNING ON A LITTLE ENDIAN MACHINE */
  for(size_t i = 0; i < max; i++) 
  {
    /* Just move the ruddy bytes around */
    data[seedLen  ]  = (i >> 24) & 0xFF;
    data[seedLen+1]  = (i >> 16) & 0xFF;
    data[seedLen+2]  = (i >>  8) & 0xFF;
    data[seedLen+3]  =  i        & 0xFF;

    SHA1(data, dataLen, tmp);
    tmp += hLen;
  }

  free(data);

  return output;
}

char *strdup(const char *src)
{
  if (src == NULL)
    return NULL;

  char *dest = calloc(strlen(src)+1, sizeof(*dest));
  strcpy(dest,src);
  
  return dest;
}
